# ------------------------------------------------------------------------------
# Copyright (c) HQU
# Licensed under the HQU License.
# Written by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# Modified by Wang Youjije (youjieWang@stu.hqu.edu.cn)
# ------------------------------------------------------------------------------
import torch
import torch.nn as nn


# 模型代码后面弄一下接口写得更好
class HRCNT(nn.Module):
    def __init__(self, cfg, depths=[3, 4, 6, 3], k_size=[3, 3, 3, 3], **kwargs):
        super(HRCNT, self).__init__()

    def forward(self):
        pass


def get_pose_net(cfg, is_train, **kwargs):
    '''
    :param cfg: cfg参数
    :param is_train: 是否是训练状态
    :param kwargs: 其他参数
    :return:返回模型
    '''
    model = HRCNT(cfg, **kwargs)

    if is_train and cfg['MODEL']['INIT_WEIGHTS']:  # 如果是训练状态并且有初始化权重就模型初始化权重
        model.init_weights(cfg['MODEL']['PRETRAINED'])

    return model
